# -*-coding:utf-8-*-

import numpy as np
from sklearn.feature_extraction import DictVectorizer
from sklearn import preprocessing
from sklearn import tree



class DicTree:
    def __init__(self):
        self.features_list = []
        data = np.loadtxt("./data/1.txt", delimiter="\t", dtype=np.str)
        x_data = data[1:, :-1]
        self.y_data = data[1:, -1:]
        title = data[:1, :]
        # DictVectorizer将dict类型的list数据，转换成numpy array
        self.vec = DictVectorizer()
        for row in x_data:
            DictRow = {}
            for f in range(0, len(row)):
                DictRow[title[0][f]] = row[f]
            self.features_list.append(DictRow)
        print(self.features_list)

    def transData(self):
        X = self.vec.fit_transform(self.features_list).toarray()
        print(X)
        lb = preprocessing.LabelBinarizer()
        y = lb.fit_transform(self.y_data)
        print("*" * 30)
        print(y)
        return X, y

    def dicsionTree(self):
        model = tree.DecisionTreeClassifier(criterion='entropy')
        X, y = d.transData()
        # 训练
        model.fit(X, y)
        train = [[1., 0., 0., 0., 1., 1., 0., 0., 0., 1., 0.]]
        p = model.predict(train)
        with open("tree.dot", 'w') as f:
            f = tree.export_graphviz(model, out_file=f)
        print(p)


d = DicTree()
d.transData()
d.dicsionTree()
